import os
import cv2
import torch
from torch.utils.data import Dataset
from torchvision import transforms
from torchvision.datasets.folder import IMG_EXTENSIONS
import utils

class DatasetLoader(Dataset):
    def __init__(self, root_dir):
        self.root_dir = root_dir
        self.samples = []

        files = os.listdir(root_dir)


        for file in files:
            if file.endswith('.png') and "_matte" not in file:
                file_name_base = file[:-4]
                matte_name = f"{file_name_base}_matte.png" # 标签label
                if matte_name in files:
                    self.samples.append(
                        (
                            os.path.join(root_dir, file),
                            os.path.join(root_dir, matte_name)
                        )
                    )

        # transform
        mean, std = utils.calc_mean_std_func(self.root_dir)
        print(f"数据集中的 mean:{mean}, std:{std}")

        # 转tensor和归一化
        self.img_transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize(mean=mean, std=std)
        ])

    def __getitem__(self, index):
        img_path, label_matte_path = self.samples[index]

        # 读取图片
        img = cv2.imread(img_path)
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = cv2.resize(img, (224, 224), interpolation=cv2.INTER_LINEAR) # 图片长度转换，彩色一般使用双线性插值

        # 读取matte图片，掩码图，标签
        label_matte = cv2.imread(label_matte_path, 0) # 灰度模式
        label_matte = cv2.resize(label_matte, (224, 224), interpolation=cv2.INTER_NEAREST) # 掩码图长度转换，使用最近邻插值

        # 处理2值化，0或者1
        label_matte[label_matte > 0] = 1

        # 转tensor和归一化
        img_tensor = self.img_transform(img)
        label_matte_tensor = torch.from_numpy(label_matte).long()

        return img_tensor, label_matte_tensor

    def __len__(self):
        return len(self.samples)



if __name__ == "__main__":
    # 测试
    dataset_path = "../Portrait-dataset-2000/dataset/training"

    if not os.path.exists(dataset_path):
        print(f"目录不存在: {dataset_path}")
    else:
        dataset = DatasetLoader(dataset_path)
        print(f"数据集长度: {len(dataset)}")
        if len(dataset) > 0:
            img, mask = dataset[0]
            print(f"第一条样本 img shape: {img.shape}, mask shape: {mask.shape}")



